/-
Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
module
prelude
public import Lean.Meta.Tactic.Grind.Types
import Init.Grind.Util
import Lean.Meta.Tactic.Grind.Proof
import Lean.Meta.Tactic.Grind.MatchCond
import Lean.Meta.Tactic.Grind.Util
namespace Lean.Meta.Grind
/-!
Debugging support code for checking basic invariants.
-/

def checkEqc (root : ENode) : GoalM Unit := do
  let mut size := 0
  let mut curr := root.self
  repeat
    size := size + 1
    -- The root of `curr` must be `root`
    assert! isSameExpr (← getRoot curr) root.self
    -- Check congruence root
    if curr.isApp then
      if let some { e } := (← get).congrTable.find? { e := curr } then
        if (← hasSameType e.getAppFn curr.getAppFn) then
          assert! isSameExpr e (← getCongrRoot curr)
      else
        assert! (← isCongrRoot curr)
    -- If the equivalence class does not have HEq proofs, then the types must be definitionally equal.
    unless root.heqProofs do
      assert! (← hasSameType curr root.self)
    -- Starting at `curr`, following the `target?` field leads to `root`.
    let mut n := curr
    repeat
      if let some target ← getTarget? n then
        n := target
      else
        break
    assert! isSameExpr n root.self
    -- Go to next element
    curr ← getNext curr
    if isSameExpr root.self curr then
      break
  -- The size of the equivalence class is correct.
  assert! root.size == size

def checkChild (e : Expr) (child : Expr) : GoalM Bool := do
  let some childRoot ← getRoot? child | return false
  return isSameExpr childRoot e

def checkMatchCondParent (e : Expr) (parent : Expr) : GoalM Bool := do
  let_expr Grind.MatchCond parent ← parent | return false
  let mut curr := parent
  repeat
    let .forallE _ d b _ := curr
      | return false
    match_expr d with
    | Eq _ lhs _ => if (← checkChild e lhs) then return true -- found it
    | HEq α lhs _ _ => if (← checkChild e α <||> checkChild e lhs) then return true -- found it
    | _ => pure ()
    curr := b
  return false

def checkParents (e : Expr) : GoalM Unit := do
  -- **Note**: We currently do not check the `funCC` case
  unless (← useFunCC e) do
    if (← isRoot e) then
      for parent in (← getParents e).elems do
        if isMatchCond parent then
          unless (← checkMatchCondParent e parent) do
            throwError "e: {e}, parent: {parent}"
          assert! (← checkMatchCondParent e parent)
        else
          let mut found := false
          -- There is an argument `arg` s.t. root of `arg` is `e`.
          for arg in parent.getAppArgs do
            if (← checkChild e arg) then
              found := true
              break
          -- Recall that we have support for `Expr.forallE` propagation. See `ForallProp.lean`.
          if let .forallE _ d b _ := parent then
            if (← checkChild e d) then
              found := true
            unless b.hasLooseBVars do
              if (← checkChild e b) then
                found := true
          unless found do
            unless (← checkChild e parent.getAppFn) do
              throwError "e: {e}, parent: {parent}"
            assert! (← checkChild e parent.getAppFn)
    else
      -- All the parents are stored in the root of the equivalence class.
      assert! (← getParents e).isEmpty

def checkPtrEqImpliesStructEq : GoalM Unit := do
  let exprs ← getExprs
  for h₁ : i in *...exprs.size do
    let e₁ := exprs[i]
    for h₂ : j in (i+1)...exprs.size do
      let e₂ := exprs[j]
      -- We don't have multiple nodes for the same expression
      assert! !isSameExpr e₁ e₂
      -- and the two expressions must not be structurally equal
      assert! !Expr.equal e₁ e₂

def checkProofs : GoalM Unit := do
  let eqcs ← getEqcs
  for eqc in eqcs do
    for a in eqc do
      for b in eqc do
        unless isSameExpr a b do
          let p ← mkEqHEqProof a b
          trace_goal[grind.debug.proofs] "{a} = {b}"
          check p
          trace_goal[grind.debug.proofs] "checked: {← inferType p}"

/-- Checks invariants if `grind.debug` is enabled. -/
public def checkInvariants (expensive := false) : GoalM Unit := do
  if grind.debug.get (← getOptions) then
    for e in (← getExprs) do
      let node ← getENode e
      checkParents node.self
      if node.isRoot then
        checkEqc node
    if expensive then
      checkPtrEqImpliesStructEq
    Solvers.checkInvariants
  if expensive && grind.debug.proofs.get (← getOptions) then
    checkProofs

@[inherit_doc Grind.checkInvariants]
public def Goal.checkInvariants (goal : Goal) (expensive := false) : GrindM Unit :=
  discard <| GoalM.run' goal <| Grind.checkInvariants expensive

end Lean.Meta.Grind
